package com.ys.util;

import java.util.*;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.ys.config.MyConfig;
import org.apache.http.HttpHost;
import org.apache.http.auth.AuthScope;
import org.apache.http.auth.UsernamePasswordCredentials;
import org.apache.http.client.CredentialsProvider;
import org.apache.http.entity.ContentType;
import org.apache.http.impl.client.BasicCredentialsProvider;
import org.apache.http.impl.nio.client.HttpAsyncClientBuilder;
import org.apache.http.nio.entity.NStringEntity;
import org.apache.http.util.EntityUtils;
import org.elasticsearch.client.Request;
import org.elasticsearch.client.Response;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.RestClientBuilder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;

/**
 * 百度Elasticsearch BES 相关
 * Elasticsearch Java客户端选择：Low Level REST Client 7.x 。
 * Low Level REST Client 7.x 相关操作通过HTTP请求与ES集群进行通信，与所有的ES版本兼容。
 * High Level REST Client 7.x 暂不支持Elasticsearch向量数据库相关操作
 */
@Component
public class RestClientUtil {
    private static final Logger logger = LoggerFactory.getLogger(RestClientUtil.class);

    //获取百度Elasticsearch参数
    private static final String BES_USERNAME = MyConfig.getBaiduElasticsearchUsername();
    private static final String BES_PASSWORD = MyConfig.getBaiduElasticsearchPassword();
    private static final String BES_HOSTNAME = MyConfig.getBaiduElasticsearchHostname();
    private static final Integer BES_PORT = Integer.parseInt(MyConfig.getBaiduElasticsearchPort());


    /**
     * 创建向量数据库索引----对应百度Embedding-V1模型
     * 成功返回索引名称，失败返回null
     */
    public String putIndex() throws Exception{
        //雪花算法获取id
        SnowflakeIdWorker idWorker = new SnowflakeIdWorker(0, 0);
        long nextId = idWorker.nextId();
        //获取入参
        String indexName = "eknow_index_"+nextId;//索引名称，唯一、非空；命名示例：eknow_index_1
        //创建向量数据库索引（Embedding-V1），固定参数
        String entity ="{\"settings\": {\"index\": {\"knn\": true,\"number_of_shards\" : 10,\"number_of_replicas\" : 1}},\"mappings\": {\"properties\": { \"content_vector\": { \"type\": \"bpack_vector\", \"dims\": 384, \"index_type\":  \"hnsw\",  \"space_type\": \"l2\", \"parameters\": { \"ef_construction\": 200, \"m\": 32 },\"engine\" : \"turbo\" }, \"shard_id\": { \"type\": \"keyword\" }, \"enabled\" : { \"type\" : \"boolean\" }, \"knowledge_id\": { \"type\": \"keyword\" } } } }";
        //创建RestClient
        RestClient restClient = getRestClient();
        // 创建request
        Request request = new Request("PUT","/"+indexName);
        // 配置请求体
        request.setEntity(new NStringEntity(entity, ContentType.APPLICATION_JSON));
        Response response = restClient.performRequest(request);
        String responseStr = EntityUtils.toString(response.getEntity());
        logger.info("创建向量数据库索引返回参数："+responseStr);
        JSONObject jsonObject = JSONObject.parseObject(responseStr);
        if((boolean) jsonObject.get("acknowledged")){
            return indexName;
        }else {
            return null;
        }
    }


    /**
     * 向量新增
     * indexName索引名称，contentVector向量数据
     * 返回新增的向量id
     */
    public String putVectorByIndex(String indexName, List<Double> contentVector) throws Exception{
        if(indexName==null || contentVector==null) return null;
        //雪花算法获取id
        SnowflakeIdWorker idWorker = new SnowflakeIdWorker(0, 0);
        long nextId = idWorker.nextId();
        String shardId = "eknow_shard_"+nextId;
        //创建RestClient
        RestClient restClient = getRestClient();
        // 创建request
        Request request = new Request("PUT","/_bulk");

        //创建向量数据库索引（Embedding-V1），固定参数
        String entity ="{ \"index\" : { \"_index\" : \""+indexName+"\", \"_id\" : \""+shardId+"\"} }";

        Map<String ,Object> params = new HashMap<>();
        params.put("content_vector",contentVector);//向量数据
        params.put("enabled",true);//是否启用该分段数据
        params.put("shard_id",shardId);//该分类碎片的id--预留字段，暂时和该向量数据id存同一个
        params.put("knowledge_id",indexName);//知识库的id--预留字段，暂时和该知识库的id存同一个
        //参数格式转换
        ObjectMapper objectMapper = new ObjectMapper();
        String sentStr = objectMapper.writeValueAsString(params);
        entity +="\n" + sentStr +"\n";
        logger.info("向量新增入参："+entity);
        // 配置请求体
        request.setEntity(new NStringEntity(entity, ContentType.APPLICATION_JSON));
        Response response = restClient.performRequest(request);
        String responseStr = EntityUtils.toString(response.getEntity());
        logger.info("向量新增返回参数："+responseStr);
        //获取回参
        //回参参考数据：{"took":5,"errors":false,"items":[{"index":{"_index":"eknow_index_1187080052005666816","_type":"_doc","_id":"eknow_shard_1187085817735020544","_version":1,"result":"created","_shards":{"total":2,"successful":2,"failed":0},"_seq_no":0,"_primary_term":1,"status":201}}]}
        JSONObject jsonObject = JSONObject.parseObject(responseStr);
        if(!(boolean) jsonObject.get("errors")){
            return shardId;
        }else {
            return null;
        }

    }

    /**
     * 向量查询，根据置信度倒序
     * indexNameList索引名称集合（多个索引）、contentVector向量数据、size查询数量、score置信度
     * 返回对应id集合
     */
    public List<Map<String,Object>> searchVectorByIndex(List<String> indexNameList, List<Double> contentVector,int size,float score) throws Exception{
        if(indexNameList==null || contentVector==null) return null;
        //返回参数
        List<Map<String,Object>> resultList = new ArrayList<>();
        String indexNameListStr = "";
        if(indexNameList.size()>0){
            for(String item:indexNameList){
                if(indexNameListStr.equals("")) {
                    indexNameListStr = item;
                }else {
                    indexNameListStr += ","+item;
                }
            }
        }
        //创建RestClient
        RestClient restClient = getRestClient();
        // 创建request
        Request request = new Request("GET",indexNameListStr+"/_search");
        //入参
        Map<String,Object> inputMap = new HashMap<>();
        inputMap.put("size",size);//返回数
        inputMap.put("_source",Arrays.asList("id"));//返回参数，默认id
        JSONObject sort = new JSONObject();sort.put("_score","desc");
        inputMap.put("sort",sort);//排序，根据置信度倒序
        JSONObject content = new JSONObject();
        content.put("vector",contentVector);
        content.put("k",10);
        content.put("ef",256);
        JSONObject knn = new JSONObject();
        knn.put("content_vector",content);
        JSONObject query = new JSONObject();
        query.put("knn",knn);
        inputMap.put("query",query);//查询的向量数据
        //参数转换
        ObjectMapper objectMapper = new ObjectMapper();
        String entity = objectMapper.writeValueAsString(inputMap);
        logger.info("向量查询入参："+entity);
        // 配置请求体
        request.setEntity(new NStringEntity(entity, ContentType.APPLICATION_JSON));
        Response response = restClient.performRequest(request);
        String responseStr = EntityUtils.toString(response.getEntity());
        logger.info("向量查询返回参数："+responseStr);
        //回参参考数据：{"took":3,"timed_out":false,"_shards":{"total":10,"successful":10,"skipped":0,"failed":0},"hits":{"total":{"value":1,"relation":"eq"},"max_score":0.4689896,"hits":[{"_index":"eknow_index_1187080052005666816","_type":"_doc","_id":"eknow_shard_1187083202859106304","_score":0.4689896,"_source":{}}]}}
        //获取回参
        JSONObject jsonObject = JSONObject.parseObject(responseStr);
        if(jsonObject.get("hits")!=null){
            JSONObject hits = JSONObject.parseObject(jsonObject.get("hits").toString());
            JSONArray hitList = hits.getJSONArray("hits");
            if(hitList!=null && hitList.size()>0){
                for(int i=0;i<hitList.size();i++){
                    Map<String,Object> info = new HashMap<>();
                    JSONObject msg = (JSONObject)hitList.get(i);
                    float backScore = Float.parseFloat(msg.get("_score").toString());
                    if(backScore-score>=0){
                        info.put("score",msg.get("_score"));//置信度
                        info.put("id",msg.get("_id"));//id
                        resultList.add(info);
                    }
                }
            }
            return resultList;
        }else {
            return null;
        }
    }


    /**
     * 向量批量新增
     * indexName索引名称，contentVector向量数据
     * 返回新增的向量id
     */
    public List<String> putVectorByIndexByList(String indexName,  List<JSONObject> embeddingList) throws Exception{
        if(indexName==null || embeddingList==null) return null;
        //雪花算法获取id
        SnowflakeIdWorker idWorker = new SnowflakeIdWorker(0, 0);

        //创建RestClient
        RestClient restClient = getRestClient();
        // 创建request
        Request request = new Request("PUT","/_bulk");

        List<String> shardIdList = new ArrayList<>();
        String entity = "";
        if(embeddingList!=null && embeddingList.size()>0){
            for(JSONObject item : embeddingList){
                long nextId = idWorker.nextId();
                String shardId = "eknow_shard_"+nextId;
                List<Double> embedding = (List<Double>)item.get("embedding");//向量内容
                //创建向量数据库索引（Embedding-V1），固定参数
                entity +="{ \"index\" : { \"_index\" : \""+indexName+"\", \"_id\" : \""+shardId+"\"} }";

                Map<String ,Object> params = new HashMap<>();
                params.put("content_vector",embedding);//向量数据
                params.put("enabled",true);//是否启用该分段数据
                params.put("shard_id",shardId);//该分类碎片的id--预留字段，暂时和该向量数据id存同一个
                params.put("knowledge_id",indexName);//知识库的id--预留字段，暂时和该知识库的id存同一个
                //参数格式转换
                ObjectMapper objectMapper = new ObjectMapper();
                String sentStr = objectMapper.writeValueAsString(params);
                entity +="\n" + sentStr +"\n";
                shardIdList.add(shardId);
            }
        }

        logger.info("向量新增入参："+entity);
        // 配置请求体
        request.setEntity(new NStringEntity(entity, ContentType.APPLICATION_JSON));
        Response response = restClient.performRequest(request);
        String responseStr = EntityUtils.toString(response.getEntity());
        logger.info("向量新增返回参数："+responseStr);
        //获取回参
        //回参参考数据：{"took":5,"errors":false,"items":[{"index":{"_index":"eknow_index_1187080052005666816","_type":"_doc","_id":"eknow_shard_1187085817735020544","_version":1,"result":"created","_shards":{"total":2,"successful":2,"failed":0},"_seq_no":0,"_primary_term":1,"status":201}}]}
        JSONObject jsonObject = JSONObject.parseObject(responseStr);
        if(!(boolean) jsonObject.get("errors")){
            return shardIdList;
        }else {
            return null;
        }

    }



    /**
     * 向量批量新增,外加参数
     * indexName索引名称，contentVector向量数据
     * 返回新增的向量id
     */
    public List<String> putVectorWithDataByList(String indexName,  List<JSONObject> embeddingList,String dataShardId,String knowledgeId) throws Exception{
        if(indexName==null || embeddingList==null) return null;
        //雪花算法获取id
        SnowflakeIdWorker idWorker = new SnowflakeIdWorker(0, 0);

        //创建RestClient
        RestClient restClient = getRestClient();
        // 创建request
        Request request = new Request("PUT","/_bulk");

        List<String> shardIdList = new ArrayList<>();
        String entity = "";
        if(embeddingList!=null && embeddingList.size()>0){
            for(JSONObject item : embeddingList){
                long nextId = idWorker.nextId();
                String shardId = "eknow_shard_"+nextId;
                List<Double> embedding = (List<Double>)item.get("embedding");//向量内容
                //创建向量数据库索引（Embedding-V1），固定参数
                entity +="{ \"index\" : { \"_index\" : \""+indexName+"\", \"_id\" : \""+shardId+"\"} }";

                Map<String ,Object> params = new HashMap<>();
                params.put("content_vector",embedding);//向量数据
                params.put("enabled",true);//是否启用该分段数据
                params.put("shard_id",dataShardId);//该分类碎片的id--预留字段，暂时和该向量数据id存同一个
                params.put("knowledge_id",knowledgeId);//知识库的id--预留字段，暂时和该知识库的id存同一个
                //参数格式转换
                ObjectMapper objectMapper = new ObjectMapper();
                String sentStr = objectMapper.writeValueAsString(params);
                entity +="\n" + sentStr +"\n";
                shardIdList.add(shardId);
            }
        }

        logger.info("向量新增入参："+entity);
        // 配置请求体
        request.setEntity(new NStringEntity(entity, ContentType.APPLICATION_JSON));
        Response response = restClient.performRequest(request);
        String responseStr = EntityUtils.toString(response.getEntity());
        logger.info("向量新增返回参数："+responseStr);
        //获取回参
        //回参参考数据：{"took":5,"errors":false,"items":[{"index":{"_index":"eknow_index_1187080052005666816","_type":"_doc","_id":"eknow_shard_1187085817735020544","_version":1,"result":"created","_shards":{"total":2,"successful":2,"failed":0},"_seq_no":0,"_primary_term":1,"status":201}}]}
        JSONObject jsonObject = JSONObject.parseObject(responseStr);
        if(!(boolean) jsonObject.get("errors")){
            return shardIdList;
        }else {
            return null;
        }

    }

    /**
     * 向量修改
     * 成功返回数据Id，失败返回null
     */
    public String updateVectorByIndex(String indexName,String dataId, List<Double> contentVector,boolean enabled) throws Exception{
        if(indexName==null || dataId==null || contentVector==null) return null;

        //创建RestClient
        RestClient restClient = getRestClient();
        // 创建request
        Request request = new Request("POST","/"+indexName+"/_update/"+dataId);
        //入参
        Map<String,Object> inputMap = new HashMap<>();
        JSONObject content = new JSONObject();
        content.put("content_vector",contentVector);
        content.put("enabled",enabled);
        inputMap.put("doc",content);
        //参数转换
        ObjectMapper objectMapper = new ObjectMapper();
        String entity = objectMapper.writeValueAsString(inputMap);
        logger.info("向量修改入参："+entity);
        request.setEntity(new NStringEntity(entity, ContentType.APPLICATION_JSON));

        Response response = restClient.performRequest(request);
        String responseStr = EntityUtils.toString(response.getEntity());
        logger.info("修改向量数据返回参数："+responseStr);
        //获取回参
        //回参参考数据：{"_index":"eknow_index_1187080052005666816","_type":"_doc","_id":"eknow_shard_1187083202859106304","_version":2,"result":"updated","_shards":{"total":2,"successful":2,"failed":0},"_seq_no":1,"_primary_term":1}
        JSONObject jsonObject = JSONObject.parseObject(responseStr);
        if(jsonObject.get("result")!=null && jsonObject.get("result").equals("updated")){
            return dataId;
        }else {
            return null;
        }
    }



    /**
     * 向量删除
     * indexName索引名称，dataId要删除的数据Id
     * * 成功返回数据Id，失败返回null
     */
    public String deleteVectorByIndex(String indexName,String dataId) throws Exception{
        if(indexName==null || dataId==null ) return null;
        //创建RestClient
        RestClient restClient = getRestClient();
        // 创建request
        Request request = new Request("DELETE","/"+indexName+"/_doc/"+dataId);
        Response response = restClient.performRequest(request);
        String responseStr = EntityUtils.toString(response.getEntity());
        logger.info("删除向量数据返回参数："+responseStr);
        //获取回参
        //回参参考数据：{"_index":"eknow_index_1187080052005666816","_type":"_doc","_id":"eknow_shard_1187083202859106304","_version":3,"result":"deleted","_shards":{"total":2,"successful":2,"failed":0},"_seq_no":2,"_primary_term":1}
        JSONObject jsonObject = JSONObject.parseObject(responseStr);
        if(jsonObject.get("result")!=null && jsonObject.get("result").equals("deleted")){
            return dataId;
        }else {
            return null;
        }
    }

    /**
     * 删除向量数据库索引----类似删库，慎用
     * indexName索引名称
     * 成功返回索引名称，失败返回null
     */
    public String deletelIndex(String indexName) throws Exception{
        if(indexName==null) return null;
        //创建RestClient
        RestClient restClient = getRestClient();
        // 创建request
        Request request = new Request("DELETE","/"+indexName);
        Response response = restClient.performRequest(request);
        String responseStr = EntityUtils.toString(response.getEntity());
        logger.info("删除向量数据库索引返回参数："+responseStr);
        JSONObject jsonObject = JSONObject.parseObject(responseStr);
        if((boolean) jsonObject.get("acknowledged")){
            return indexName;
        }else {
            return null;
        }
    }



    //百度es创建RestClient公共方法
    public RestClient getRestClient(){
        final CredentialsProvider credentialsProvider = new BasicCredentialsProvider();
        //访问用户名和密码为您创建Elasticsearch实例时设置的用户名和密码。
        credentialsProvider.setCredentials(AuthScope.ANY, new UsernamePasswordCredentials(BES_USERNAME, BES_PASSWORD));
        // 通过builder创建rest client，配置http client的HttpClientConfigCallback。
        // ES HTTP URL 在Baidu Elasticsearch界面中可以查询
        RestClient restClient = RestClient.builder(new HttpHost(BES_HOSTNAME, BES_PORT))
                .setHttpClientConfigCallback(new RestClientBuilder.HttpClientConfigCallback() {
                    @Override
                    public HttpAsyncClientBuilder customizeHttpClient(HttpAsyncClientBuilder httpClientBuilder) {
                        return httpClientBuilder.setDefaultCredentialsProvider(credentialsProvider);
                    }
                }).build();
        return restClient;
    }


}
